import socket
import struct

# QKD Packet Fields
HEADER_FORMAT = "!I I B B I I 32s"

class QKDAppPacket:
    def __init__(self, length, msg_id, encryption_algo, auth_algo, encryption_key_id, auth_key_id, auth_tag, payload):
        self.length = length
        self.msg_id = msg_id
        self.encryption_algo = encryption_algo
        self.auth_algo = auth_algo
        self.encryption_key_id = encryption_key_id
        self.auth_key_id = auth_key_id
        self.auth_tag = auth_tag
        self.payload = payload

    def pack(self):
        # Pack the header and payload into bytes
        header = struct.pack(HEADER_FORMAT, self.length, self.msg_id, self.encryption_algo, self.auth_algo,
                            self.encryption_key_id, self.auth_key_id, self.auth_tag)
        return header + self.payload

    @staticmethod
    def unpack(data):
        # Unpack the received data into the QKDAppPacket fields
        header_size = struct.calcsize(HEADER_FORMAT)
        header = data[:header_size]
        payload = data[header_size:]
        unpacked_data = struct.unpack(HEADER_FORMAT, header)
        return QKDAppPacket(*unpacked_data, payload)

# Sender Code
def sender():
    sender_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    receiver_address = ('127.0.0.1', 12345)

    # Packet fields
    length = 100
    msg_id = 1
    encryption_algo = 1  # Example encryption algorithm ID
    auth_algo = 1  # Example authentication algorithm ID
    encryption_key_id = 101
    auth_key_id = 201
    auth_tag = b'auth_tag_example'.ljust(32, b'\0')  # Pad to 32 bytes
    payload = b'This is the protected payload'

    # Create packet
    packet = QKDAppPacket(length, msg_id, encryption_algo, auth_algo, encryption_key_id, auth_key_id, auth_tag, payload)
    packed_data = packet.pack()

    # Send packet
    sender_socket.sendto(packed_data, receiver_address)
    print("Packet sent")
    sender_socket.close()

# Receiver Code
def receiver():
    receiver_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    receiver_socket.bind(('127.0.0.1', 12345))
    print("Receiver is listening...")

    # Receive packet
    data, addr = receiver_socket.recvfrom(4096)
    print("Packet received from", addr)

    # Unpack packet
    packet = QKDAppPacket.unpack(data)
    print(f"Received Packet:\n Length: {packet.length}\n Msg-ID: {packet.msg_id}\n Encryption Algorithm: {packet.encryption_algo}\n Authentication Algorithm: {packet.auth_algo}\n Encryption Key ID: {packet.encryption_key_id}\n Authentication Key ID: {packet.auth_key_id}\n Authentication Tag: {packet.auth_tag}\n Payload: {packet.payload.decode('utf-8')}")
    
    receiver_socket.close()

if __name__ == "__main__":
    # For testing, run sender and receiver in separate processes or threads
    import threading

    receiver_thread = threading.Thread(target=receiver)
    sender_thread = threading.Thread(target=sender)

    receiver_thread.start()
    sender_thread.start()

    sender_thread.join()
    receiver_thread.join()
